"""模型预测模块."""

from __future__ import annotations

import json
import logging
from typing import Optional

import pandas as pd
from pydantic import BaseModel, Field

from ..models.storage import ModelLoader, ModelStorage
from ..utils.state import GLOBAL_DATA_STORE
from ..utils import ETFDataTool, FeatureEngineeringTool
from ..quant import QuantAnalyzer

logger = logging.getLogger(__name__)


class ModelPredictionInput(BaseModel):
    """模型预测输入模型."""

    symbol: str = Field(..., description="ETF 标准化代码，例如 159919.SZ")
    model_name: Optional[str] = Field(
        None, description="模型名称（可选），如果为 None 则使用最佳模型"
    )


def format_prediction_text(result: dict) -> str:
    """将预测结果格式化为易读的中文文本，便于大模型理解."""
    lines = []
    
    # 标题
    lines.append(f"【{result['symbol']} ETF 模型预测结果】")
    lines.append("")
    
    # 预测信息
    lines.append("🔮 预测结果：")
    prediction = result["prediction"]
    proba = result["up_probability"]
    confidence = result["confidence"]
    
    # 根据预测方向显示不同的表情符号
    if prediction == "上涨":
        trend_emoji = "📈"
    else:
        trend_emoji = "📉"
    
    lines.append(f"  • 预测方向: {trend_emoji} {prediction}")
    lines.append(f"  • 上涨概率: {proba:.2%}")
    lines.append(f"  • 置信度: {confidence:.2%}")
    lines.append(f"  • 预测日期: {result['prediction_date']}")
    lines.append("")
    
    # 模型信息
    lines.append("🤖 模型信息：")
    lines.append(f"  • 模型名称: {result['model_name']}")
    metadata = result.get("model_metadata", {})
    if metadata.get("training_date"):
        lines.append(f"  • 训练日期: {metadata['training_date']}")
    
    # 性能指标
    if metadata.get("performance_metrics"):
        metrics = metadata["performance_metrics"]
        lines.append("  • 模型性能指标：")
        if "sharpe_ratio" in metrics:
            lines.append(f"    - 夏普比率: {metrics['sharpe_ratio']:.4f}")
        if "max_drawdown" in metrics:
            lines.append(f"    - 最大回撤: {metrics['max_drawdown']:.4f}")
        if "win_rate" in metrics:
            lines.append(f"    - 胜率: {metrics['win_rate']:.2%}")
    
    lines.append("")
    
    # 使用提示
    lines.append("💡 提示：")
    lines.append("  • 预测基于历史数据训练的机器学习模型")
    lines.append("  • 预测结果仅供参考，不构成投资建议")
    lines.append("  • 建议结合风险分析和其他工具进行综合判断")
    
    return "\n".join(lines)


class ModelPrediction:
    """模型预测功能类."""

    def __init__(self, *, model_storage_dir: str = "models") -> None:
        """初始化模型预测器.
        
        Parameters
        ----------
        model_storage_dir:
            模型存储目录，默认为 "models"
        """
        self.storage = ModelStorage(model_storage_dir)
        self.loader = ModelLoader(self.storage)

    def get_json(self, symbol: str, model_name: Optional[str] = None) -> str:
        """返回 JSON 格式的预测结果.

        Parameters
        ----------
        symbol:
            ETF 标准化代码
        model_name:
            模型名称（可选），如果为 None 则使用最佳模型

        Returns
        -------
        JSON 格式的预测结果
        """
        logger.info(f"开始模型预测: symbol={symbol}, model_name={model_name}")

        # 1. 加载模型
        if model_name is None:
            model, metadata = self.loader.load_best_model(symbol)
            model_name = metadata.model_name
        else:
            model, metadata = self.storage.load(symbol, model_name)

        # 2. 获取最新特征数据
        features_df = GLOBAL_DATA_STORE.get_features(symbol)
        if features_df is None:
            raise RuntimeError(f"未找到缓存的特征数据: {symbol}，请先调用 etf_feature_engineering_tool")

        # 3. 准备预测数据（移除目标标签列，如果存在）
        # 确保特征数据与训练时一致（不包含 target_up 列）
        if "target_up" in features_df.columns:
            X_for_prediction = features_df.drop(columns=["target_up"])
        else:
            X_for_prediction = features_df

        # 4. 对齐特征：确保预测时使用的特征与训练时一致
        # 使用模型保存的特征名称来对齐特征顺序和列
        if model.feature_names is not None and len(model.feature_names) > 0:
            # 检查特征是否匹配
            missing_features = set(model.feature_names) - set(X_for_prediction.columns)
            if missing_features:
                raise RuntimeError(
                    f"特征不匹配：模型期望的特征 {missing_features} 在数据中不存在。"
                    f"模型特征数: {len(model.feature_names)}, 数据特征数: {len(X_for_prediction.columns)}"
                )
            # 按照模型训练时的特征顺序排列
            X_for_prediction = X_for_prediction[model.feature_names]
        else:
            # 如果没有保存特征名称，记录警告但继续
            logger.warning(f"模型 {model_name} 未保存特征名称，使用原始特征顺序")

        # 使用最新数据预测
        # 检测模型类型：序列模型（LSTM、GRU、CNN）需要多行数据
        is_deep_learning_model = False
        sequence_len = 1  # 传统机器学习模型只需要1行
        
        # 检查是否是深度学习序列模型
        model_name_lower = model_name.lower()
        if any(dl_name in model_name_lower for dl_name in ["lstm", "gru", "cnn"]):
            is_deep_learning_model = True
            sequence_len = 30  # 默认序列长度
            
            # 尝试从模型配置中获取实际的序列长度
            if hasattr(model, "config") and hasattr(model.config, "params"):
                sequence_len = model.config.params.get("sequence_len", 30)
            elif hasattr(model, "config") and isinstance(model.config, dict):
                sequence_len = model.config.get("params", {}).get("sequence_len", 30)
        
        # 根据模型类型选择预测数据
        if is_deep_learning_model:
            # 序列模型：使用最后 sequence_len 行数据（滑动窗口）
            if len(X_for_prediction) < sequence_len:
                raise ValueError(
                    f"特征数据不足：序列模型需要至少 {sequence_len} 行数据，"
                    f"但只有 {len(X_for_prediction)} 行。请确保有足够的历史数据。"
                )
            X_latest = X_for_prediction.iloc[-sequence_len:]
            logger.debug(f"使用序列模型预测，数据窗口大小: {sequence_len} 行")
        else:
            # 传统机器学习模型：只使用最后一行（最新数据）
            X_latest = X_for_prediction.iloc[-1:]
            logger.debug("使用传统机器学习模型预测，数据: 1 行")

        # 预测上涨概率
        proba_array = model.predict_proba(X_latest)
        
        # 对于序列模型，predict_proba 可能返回多个预测值（每个序列一个）
        # 我们取最后一个预测值（对应最新日期）
        if is_deep_learning_model and len(proba_array) > 1:
            proba = float(proba_array[-1])
        elif len(proba_array) > 0:
            proba = float(proba_array[0])
        else:
            raise RuntimeError("模型预测结果为空")

        # 4. 构建预测结果
        # 预测日期：使用最新一行的日期（对于序列模型，X_latest有多行，取最后一行）
        prediction_date = X_for_prediction.index[-1].isoformat()
        
        result = {
            "symbol": symbol,
            "model_name": model_name,
            "prediction_date": prediction_date,
            "up_probability": float(proba),
            "prediction": "上涨" if proba >= 0.5 else "下跌",
            "confidence": float(max(proba, 1 - proba)),  # 置信度：取概率和(1-概率)的较大值，反映模型对预测的确定性
            "model_metadata": {
                "training_date": metadata.training_date,
                "performance_metrics": metadata.performance_metrics,
            },
        }

        logger.info(f"预测完成: symbol={symbol}, proba={proba:.4f}")
        return json.dumps(result, ensure_ascii=False, indent=2)

    def format_text(self, json_str: str) -> str:
        """将 JSON 格式化为易读的中文文本."""
        result = json.loads(json_str)
        return format_prediction_text(result)

    def get_formatted_text(self, symbol: str, model_name: Optional[str] = None) -> str:
        """获取格式化的中文文本."""
        json_str = self.get_json(symbol, model_name)
        return self.format_text(json_str)


# 向后兼容的别名类
class ModelPredictionTool:
    """向后兼容的别名类，实际功能由 ModelPrediction 提供."""

    def __init__(self, *, model_storage_dir: str = "models") -> None:
        """初始化工具实例."""
        self._prediction = ModelPrediction(model_storage_dir=model_storage_dir)
        # 保持向后兼容的属性
        self.storage = self._prediction.storage
        self.loader = self._prediction.loader

    def _get_json(self, symbol: str, model_name: Optional[str] = None) -> str:
        """返回 JSON 格式的预测结果（向后兼容方法）。"""
        return self._prediction.get_json(symbol, model_name)

    def _format_text(self, json_str: str) -> str:
        """将 JSON 格式化为易读的中文文本（向后兼容方法）。"""
        return self._prediction.format_text(json_str)

    def run(self, input_data: dict) -> str:
        """运行预测并返回格式化的文本（向后兼容方法，用于 LangChain 工具接口）。"""
        symbol = input_data["symbol"]
        model_name = input_data.get("model_name")
        return self._prediction.get_formatted_text(symbol, model_name)


def main() -> None:
    """测试模型预测的主函数."""
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    )
    
    print("=" * 80)
    print("模型预测工具测试")
    print("=" * 80)
    
    # 使用真实数据：需要先准备数据、特征和模型
    print("\n1. 准备测试数据...")
    test_symbol = "159796"
    data_tool = ETFDataTool()
    
    try:
        # 1.1 获取 ETF 数据
        print("   1.1 调用 etf_data_tool 获取 ETF 数据...")
        result_text = data_tool._run(symbol=test_symbol, limit=200)
        output = data_tool._get_output(symbol=test_symbol, limit=200)
        symbol = output.normalized_symbol
        
        print(f"   ✓ 成功获取 ETF 数据")
        print(f"   原始代码: {test_symbol}")
        print(f"   标准化代码: {symbol}")
        print(f"   数据条数: {output.actual_count}")
        
        # 验证数据已缓存
        dataset = GLOBAL_DATA_STORE.get_dataset(symbol)
        if dataset is None:
            print("   ✗ 警告: 数据未正确缓存到全局存储")
            return
        else:
            print(f"   ✓ 数据已缓存到全局存储")
        
        # 1.2 生成特征
        print("\n   1.2 调用 etf_feature_engineering_tool 生成特征...")
        feature_tool = FeatureEngineeringTool()
        feature_result_text = feature_tool._run(symbol=symbol, include_target=True)
        feature_output = feature_tool._get_output(symbol=symbol, include_target=True)
        
        print(f"   ✓ 特征工程完成")
        print(f"   特征行数: {feature_output['feature_rows']}")
        print(f"   特征列数: {len(feature_output['feature_columns'])}")
        
        # 验证特征已缓存
        features = GLOBAL_DATA_STORE.get_features(symbol)
        if features is None:
            print("   ✗ 警告: 特征数据未正确缓存")
            return
        else:
            print(f"   ✓ 特征数据已缓存")
        
        # 1.3 训练模型（如果还没有模型）
        print("\n   1.3 检查是否有已训练的模型...")
        prediction = ModelPrediction()
        
        # 尝试加载最佳模型
        try:
            model, metadata = prediction.loader.load_best_model(symbol)
            print(f"   ✓ 找到已训练的模型: {metadata.model_name}")
            print(f"   训练日期: {metadata.training_date}")
            has_model = True
        except (ValueError, FileNotFoundError):
            print("   ⚠ 未找到已训练的模型，开始训练...")
            print("   注意：模型训练可能需要较长时间...")
            
            quant_analyzer = QuantAnalyzer()
            try:
                quant_output = quant_analyzer.analyze(symbol=symbol, include_lstm=False, enable_optimization=False)
                quant_result_text = quant_analyzer.format_result(quant_output)
                quant_data = json.loads(quant_output)
                
                print(f"   ✓ 模型训练完成")
                print(f"   最佳模型: {quant_data.get('best_model', 'N/A')}")
                has_model = True
            except Exception as e:
                print(f"   ✗ 模型训练失败: {e}")
                print("   提示：可以使用已训练的模型进行测试，或先运行 quant_analysis_tool")
                import traceback
                traceback.print_exc()
                return
        
        if not has_model:
            print("   ✗ 无法继续测试：缺少已训练的模型")
            return
        
    except Exception as e:
        print(f"   ✗ 准备测试数据失败: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # 测试1: 使用最佳模型预测
    print("\n2. 测试模型预测（使用最佳模型）...")
    try:
        # 调用工具，获取格式化的中文文本
        result_text = prediction.get_formatted_text(symbol=symbol)
        
        # 显示格式化文本
        print("返回的格式化文本：")
        print("-" * 80)
        print(result_text)
        print("-" * 80)
        
        # 通过 get_json 方法获取原始 JSON 数据，用于验证
        result_json = prediction.get_json(symbol=symbol)
        result_data = json.loads(result_json)
        
        print(f"\n   ✓ 预测成功")
        print(f"   ETF 代码: {result_data['symbol']}")
        print(f"   模型名称: {result_data['model_name']}")
        print(f"   预测方向: {result_data['prediction']}")
        print(f"   上涨概率: {result_data['up_probability']:.4f} ({result_data['up_probability']:.2%})")
        print(f"   置信度: {result_data['confidence']:.4f} ({result_data['confidence']:.2%})")
        print(f"   预测日期: {result_data['prediction_date']}")
        
        if result_data.get('model_metadata'):
            metadata = result_data['model_metadata']
            if metadata.get('training_date'):
                print(f"   模型训练日期: {metadata['training_date']}")
            if metadata.get('performance_metrics'):
                metrics = metadata['performance_metrics']
                if 'sharpe_ratio' in metrics:
                    print(f"   模型夏普比率: {metrics['sharpe_ratio']:.4f}")
        
    except Exception as e:
        print(f"   ✗ 测试失败: {e}")
        import traceback
        traceback.print_exc()
        return
    
    # 测试2: 指定模型名称预测（如果可能）
    print("\n3. 测试指定模型名称预测...")
    try:
        # 先获取可用的模型列表
        models = prediction.storage.list_models(symbol)
        if models:
            # 使用第一个可用模型
            test_model_name = models[0]["model_name"]
            print(f"   使用模型: {test_model_name}")
            
            result_text = prediction.get_formatted_text(symbol=symbol, model_name=test_model_name)
            result_json = prediction.get_json(symbol=symbol, model_name=test_model_name)
            result_data = json.loads(result_json)
            
            print(f"   ✓ 指定模型预测成功")
            print(f"   使用的模型: {result_data['model_name']}")
            print(f"   预测方向: {result_data['prediction']}")
            print(f"   上涨概率: {result_data['up_probability']:.2%}")
        else:
            print("   ⚠ 没有可用的模型进行测试")
    except Exception as e:
        print(f"   ⚠ 指定模型预测测试跳过: {e}")
    
    # 测试3: 错误处理
    print("\n4. 测试错误处理...")
    try:
        prediction.get_json(symbol="NONEXISTENT.SZ")
        print("   ✗ 应该抛出异常但没有")
    except RuntimeError as e:
        print(f"   ✓ 正确抛出异常: {e}")
    except Exception as e:
        print(f"   ⚠ 抛出其他异常: {type(e).__name__}: {e}")
    
    print("\n" + "=" * 80)
    print("测试完成！")
    print("=" * 80)


if __name__ == "__main__":
    import sys
    
    # 检查是否作为模块运行
    if not __package__:
        print("错误: 请使用以下命令运行测试:")
        print("  python -m src.workflows.model_prediction")
        print("或者:")
        print("  uv run python -m src.workflows.model_prediction")
        sys.exit(1)
    
    main()

